from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from langchain.schema.runnable import RunnableLambda
from langserve import add_routes
from patagent.task import translation_chain
from patagent.task import quality_chain
from dotenv import load_dotenv
from patagent.task.fewshot_extractor import fewshot_extractor
from patagent.agent import qa_graph

load_dotenv()

app = FastAPI(
    title="Patsnap Agent Server",
    version="1.0",
    description="A agents api server using Langchain's Runnable interfaces",
)

# Set all CORS enabled origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],
)

def rag(question: str) -> str:
    try:
        res = qa_graph.invoke({'question': question})
        return res['generation']
    except Exception as e:
        print(f"Raise Exception: {e}")
        return '您的问题超出了我的认知，请重新尝试'

# RAG
add_routes(
    app,
    RunnableLambda(rag),
    path="/rag",
)

# merterial translate
add_routes(
    app,
    translation_chain,
    path="/translate",
)

# merterial translate
add_routes(
    app,
    quality_chain,
    path="/quality",
)

# image ocr
add_routes(
    app,
    fewshot_extractor,
    path="/extractor",
)

def main():
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=1888)

if __name__ == "__main__":
    main()